#include <stdio.h>
#include "rc6.h"
#include <malloc.h>
#include <string.h>
#include <math.h>
#ifdef EN_RC6_BASE64_ENCODE
#include "base64.h"
#endif
    
// #define w 32        /* word 明文的字长 bits */
// #define r 20        /* round 加密轮数 */


#define P32 0xB7E15163      /* 定义两个常量，用于子密钥生成 */
#define Q32 0x9E3779B9

#define bytes   (ctx->opt.w / 8)                        /* 定义字节*/
#define c       ((ctx->opt.b + bytes - 1) / bytes)      /* 密钥字数 */
#define R24     (2 * ctx->opt.r + 4)
#define lgw     5                           /* log2(w)  */

/* 定义逻辑移位操作运算 */
#define ROTL(x,y) (((x)<<(y&(ctx->opt.w-1))) | ((x)>>(ctx->opt.w-(y&(ctx->opt.w-1)))))
#define ROTR(x,y) (((x)>>(y&(ctx->opt.w-1))) | ((x)<<(ctx->opt.w-(y&(ctx->opt.w-1)))))

// unsigned int S[R24 - 1];                    /* 子密钥组 */

static void rc6_key_setup(struct rc6_ctx *ctx, unsigned char *K, int b) /* b 秘钥长度 单位：字节*/
{
    int i, j, s, v;
    unsigned int L[c]; /* Big enough for max b */
    unsigned int A, B;

    L[c - 1] = 0;
    for (i = b - 1; i >= 0; i--)
        L[i / bytes] = (L[i / bytes] << 8) + K[i];

    ctx->opt.S[0] = P32;
    for (i = 1; i <= 2 * ctx->opt.r + 3; i++)
        ctx->opt.S[i] = ctx->opt.S[i - 1] + Q32;

    A = B = i = j = 0;
    v = R24;
    if (c > v) v = c;
    v *= 3;

    for (s = 1; s <= v; s++)
    {
        A = ctx->opt.S[i] = ROTL(ctx->opt.S[i] + A + B, 3);
        B = L[j] = ROTL(L[j] + A + B, A + B);
        i = (i + 1) % R24;
        j = (j + 1) % c;
    }
}

static void rc6_block_encrypt(struct rc6_ctx *ctx, unsigned int *pt, unsigned int *ct)
{
    unsigned int A, B, C, D, t, u, x;
    int i;
    int l = (int)log2(ctx->opt.w);

    A = pt[0];
    B = pt[1];
    C = pt[2];
    D = pt[3];
    B += ctx->opt.S[0];
    D += ctx->opt.S[1];
    for (i = 2; i <= 2 * ctx->opt.r; i += 2)
    {
        t = ROTL(B * (2 * B + 1), l);
        u = ROTL(D * (2 * D + 1), l);
        A = ROTL(A ^ t, u) + ctx->opt.S[i];
        C = ROTL(C ^ u, t) + ctx->opt.S[i + 1];
        x = A;
        A = B;
        B = C;
        C = D;
        D = x;
    }
    A += ctx->opt.S[2 * ctx->opt.r + 2];
    C += ctx->opt.S[2 * ctx->opt.r + 3];
    ct[0] = A;
    ct[1] = B;
    ct[2] = C;
    ct[3] = D;
}

static void rc6_block_decrypt(struct rc6_ctx *ctx, unsigned int *ct, unsigned int *pt)
{
    unsigned int A, B, C, D, t, u, x;
    int i;

    A = ct[0];
    B = ct[1];
    C = ct[2];
    D = ct[3];
    C -= ctx->opt.S[2 * ctx->opt.r + 3];
    A -= ctx->opt.S[2 * ctx->opt.r + 2];
    for (i = 2 * ctx->opt.r; i >= 2; i -= 2)
    {
        x = D;
        D = C;
        C = B;
        B = A;
        A = x;
        u = ROTL(D * (2 * D + 1), lgw);
        t = ROTL(B * (2 * B + 1), lgw);
        C = ROTR(C - ctx->opt.S[i + 1], t) ^ u;
        A = ROTR(A - ctx->opt.S[i], u) ^ t;
    }
    D -= ctx->opt.S[1];
    B -= ctx->opt.S[0];
    pt[0] = A;
    pt[1] = B;
    pt[2] = C;
    pt[3] = D;    
}

static unsigned int RC6AddPKCS7Padding(unsigned char *data, unsigned int len)
{
    unsigned int newLen;
    newLen = len + RC6_BLOCK_SIZE - (len % RC6_BLOCK_SIZE);
    memset(&data[len], newLen-len, newLen-len);
    return newLen;
}

static unsigned int RC6AddZeroPadding(unsigned char *data, unsigned int len)
{
    unsigned int newLen;
    newLen = len + RC6_BLOCK_SIZE - (len % RC6_BLOCK_SIZE);
    memset(&data[len], 0, newLen-len);
    return newLen;
}

static unsigned int RC6AddISOPadding(unsigned char *data, unsigned int len)
{
    unsigned int newLen;
    newLen = len + RC6_BLOCK_SIZE - (len % RC6_BLOCK_SIZE);
    memset(&data[len], 11, newLen-len);     
    data[newLen-1] = newLen-len;
    return newLen;
}
 
static unsigned int RC6DelPKCS7Padding(unsigned char *pData, unsigned int len)
{    
    if (0 != (len & (RC6_BLOCK_SIZE-1)) || !len) {
        return 0;
    }
    if (pData[len - 1] > len) {
        return 0;
    }
 
    return len - pData[len - 1];
}

static unsigned int RC6DelZeroPadding(unsigned char *pData, unsigned int len)
{    
    if (0 != (len & (RC6_BLOCK_SIZE-1)) || !len) {
        return 0;
    }
    while(pData[len-1]==0 && len) len--;
 
    return len;
}

static unsigned int RC6DelISOPadding(unsigned char *pData, unsigned int len)
{    
    if (0 != (len & (RC6_BLOCK_SIZE-1)) || !len) {
        return 0;
    }

    if(pData[len-1]==0 || pData[len-1]>16) return 0;

    return len - pData[len-1];
}

static unsigned int RC6AddANSIX923Padding(unsigned char *data, unsigned int len)
{
    unsigned int newLen;
    newLen = len + RC6_BLOCK_SIZE - (len % RC6_BLOCK_SIZE);
    memset(&data[len], 0, newLen-len);     
    data[newLen-1] = newLen-len;
    return newLen;
}
 
static unsigned int RC6DelANSIX923Padding(unsigned char *pData, unsigned int len)
{    
    if (0 != (len & (RC6_BLOCK_SIZE-1)) || !len) {
        return 0;
    }
    if (pData[len - 1] > len) {
        return 0;
    }
 
    return len - pData[len - 1];
}

static void XorBytes(unsigned char *pData1, const unsigned char *pData2, unsigned char nCount)
{
    unsigned char i;
    for (i = 0; i < nCount; i++) {
        pData1[i] ^= pData2[i];
    }
}

int rc6_init(struct rc6_ctx* ctx, unsigned char *key, unsigned char key_len, unsigned char *pIV, struct rc6_opt *opt)
{
    double t;
    if(!ctx || !key || key_len<8) return -1;
    if(opt)
    {
        t = log2(opt->w);
        if(t<=0 || t!=fabs(t)) return -1;   //w value is not OK 
        ctx->opt = *opt;
        if(!ctx->opt.S) 
        {
            ctx->opt.S = malloc(sizeof(int)*(2*ctx->opt.r+4));
            ctx->opt.Sm = 1;
        }
        else ctx->opt.Sm = 0;
        if(!ctx->opt.S) return -1;  //malloc fail
    }
    else    //default  w/r/b = 32/20/16
    {
        ctx->opt.w = 32;
        ctx->opt.r = 20;
        ctx->opt.S = malloc(sizeof(int)*(2*ctx->opt.r+4));
        ctx->opt.Sm = 0;
        if(!ctx->opt.S) return -1;  //malloc fail
    }
    memcpy(ctx->pIV, pIV, RC6_BLOCK_SIZE);
    ctx->opt.b = key_len;

    rc6_key_setup(ctx, key, key_len);
    return 0;
}

int rc6_encrypt(struct rc6_ctx* ctx, unsigned char *plaintext,int p_len, unsigned char *ciphertext, int *c_len)
{
    unsigned int i;
    unsigned char *pct = ciphertext;
    unsigned char pIV[RC6_BLOCK_SIZE];

    if(!plaintext || !ciphertext || !p_len) return -1;

    if (plaintext != ciphertext) {
        memcpy(ciphertext, plaintext, p_len);
    }

    //padding
    switch (ctx->padding)
    {
    case RC6_NoPadding:
        if(p_len%RC6_BLOCK_SIZE != 0) return -1;
        break;
    
    case RC6_Zeros:
        p_len = RC6AddZeroPadding(ciphertext, p_len);
        break;

    case RC6_ISO10126:
        p_len = RC6AddISOPadding(ciphertext, p_len);
        break;

    case RC6_PKCS7: 
        p_len = RC6AddPKCS7Padding(ciphertext, p_len);
        break;

    case RC6_ANSIX923:
        p_len = RC6AddANSIX923Padding(ciphertext, p_len);
        break;

    default:
        return -1;
    }

    //encrypt
    memcpy(pIV, ctx->pIV, sizeof(pIV));
    switch (ctx->mode)
    {
    case RC6_MODE_ECB:  
    case RC6_MODE_CBC:  
        for (i = p_len / RC6_BLOCK_SIZE; i > 0 ; i--, ciphertext += RC6_BLOCK_SIZE) {
            if (RC6_MODE_CBC == ctx->mode) {
                XorBytes(ciphertext, (unsigned char*)pIV, RC6_BLOCK_SIZE);
            }
            rc6_block_encrypt(ctx, (unsigned int*)ciphertext, (unsigned int*)ciphertext);
            memcpy(pIV, ciphertext, sizeof(pIV));
        }   
        break;
    case RC6_MODE_CFB:  // -- 128bit
        for (i = p_len / RC6_BLOCK_SIZE; i > 0 ; i--, ciphertext += RC6_BLOCK_SIZE) {
            rc6_block_encrypt(ctx, (unsigned int*)pIV, (unsigned int*)pIV);
            XorBytes((unsigned char*)pIV, ciphertext, RC6_BLOCK_SIZE);
            memcpy(ciphertext, pIV, RC6_BLOCK_SIZE);
        }
        break;
    case RC6_MODE_OFB:  
        for (i = p_len / RC6_BLOCK_SIZE; i > 0 ; i--, ciphertext += RC6_BLOCK_SIZE) {
            rc6_block_encrypt(ctx, (unsigned int*)pIV, (unsigned int*)pIV);
            XorBytes(ciphertext,  (unsigned char*)pIV, RC6_BLOCK_SIZE);
        }
        break;
    case RC6_MODE_CTR:
        if(ctx->ctr_alg == NULL) return -1;
        for(i = 0; i < p_len/RC6_BLOCK_SIZE; i++, ciphertext += RC6_BLOCK_SIZE){
            ctx->ctr_alg(pIV,i);
            rc6_block_encrypt(ctx, (unsigned int*)pIV, (unsigned int*)pIV);
            XorBytes(ciphertext, pIV, RC6_BLOCK_SIZE);
        }
        break;
    default:
        return -1;
    }

    #ifdef EN_RC6_BASE64_ENCODE    //using base64 encode
    unsigned char *base64 = (unsigned char *)malloc((p_len+2)*4/3);
    if(!base64) return -1;
    if(base64_encode((char*)pct, p_len, (char*)base64, &p_len) == -1)
    {
        free(base64);
        return -1;
    }
    memcpy(pct,base64,p_len);
    free(base64);
    #endif  /* EN_RC6_BASE64_ENCODE */

    if(c_len) *c_len = p_len;
    return 0;
}

int rc6_decrypt(struct rc6_ctx* ctx, unsigned char *plaintext,int *p_len, unsigned char *ciphertext, int c_len)
{
    unsigned int i;
    unsigned char *pPlainTextBack = plaintext;
    unsigned char pIV[RC6_BLOCK_SIZE] ;

    if(!plaintext || !ciphertext || !c_len) return 0;

    #ifdef EN_RC6_BASE64_ENCODE   //using base64 encode
    if(base64_decode((char*)ciphertext, c_len, (char*)plaintext, &c_len) == -1)
    {
        return 0;
    }
    #else   /* EN_RC6_BASE64_ENCODE */
    if (plaintext != ciphertext) {
        memcpy(plaintext, ciphertext, c_len);
    }
    #endif  /* EN_RC6_BASE64_ENCODE */

    //decrypt
    memcpy(pIV, ctx->pIV, sizeof(pIV));
    switch (ctx->mode)
    {
    case RC6_MODE_ECB:  
    case RC6_MODE_CBC:  
        plaintext += c_len - RC6_BLOCK_SIZE;
        for (i = c_len / RC6_BLOCK_SIZE; i > 0 ; i--, plaintext -= RC6_BLOCK_SIZE) {
            rc6_block_decrypt(ctx, (unsigned int*)plaintext, (unsigned int*)plaintext);
            if (RC6_MODE_CBC == ctx->mode) {
                if (1 == i) {
                    XorBytes(plaintext, ctx->pIV, RC6_BLOCK_SIZE);
                } else {
                    XorBytes(plaintext, plaintext - RC6_BLOCK_SIZE, RC6_BLOCK_SIZE);
                }
            }
        }
        break;
    case RC6_MODE_CFB:
        for (i = c_len / RC6_BLOCK_SIZE; i > 0 ; i--, plaintext += RC6_BLOCK_SIZE) {
            rc6_block_encrypt(ctx, (unsigned int*)pIV, (unsigned int*)pIV);
            XorBytes(pIV, plaintext, RC6_BLOCK_SIZE);
            //change
            XorBytes(pIV, plaintext, RC6_BLOCK_SIZE);
            XorBytes(plaintext, pIV, RC6_BLOCK_SIZE);
            XorBytes(pIV, plaintext, RC6_BLOCK_SIZE);
        }
        break;
    case RC6_MODE_OFB:
        for (i = c_len / RC6_BLOCK_SIZE; i > 0 ; i--, plaintext += RC6_BLOCK_SIZE) {
            rc6_block_encrypt(ctx, (unsigned int*)pIV, (unsigned int*)pIV);
            XorBytes(plaintext, (unsigned char*)pIV, RC6_BLOCK_SIZE);  
        }
        break;
    case RC6_MODE_CTR:
        if(ctx->ctr_alg == NULL) return 0;
        for(i = 0; i < c_len/RC6_BLOCK_SIZE; i++, plaintext += RC6_BLOCK_SIZE){
            ctx->ctr_alg(pIV,i);
            rc6_block_encrypt(ctx, (unsigned int*)pIV, (unsigned int*)pIV);
            XorBytes(plaintext, pIV, RC6_BLOCK_SIZE);
        }
        break;
    default:
        return 0;
    }

    //remove padding
    switch (ctx->padding)
    {
    case RC6_NoPadding:
        *p_len = c_len; 
        break;
    
    case RC6_Zeros:
        *p_len = RC6DelZeroPadding(pPlainTextBack, c_len);
        break;

    case RC6_ISO10126:
        *p_len = RC6DelISOPadding(pPlainTextBack, c_len);
        break;

    case RC6_PKCS7:
        *p_len = RC6DelPKCS7Padding(pPlainTextBack, c_len);
        break;

    case RC6_ANSIX923:
        *p_len = RC6DelANSIX923Padding(pPlainTextBack, c_len);
        break;

    default:
        return -1;
    }

    return 0;
}

int rc6_free(struct rc6_ctx* ctx)
{
    if(ctx->opt.Sm) free(ctx->opt.S);
    return 0;
}

void default_ctr_alg(void *Nonce, int counter)
{
    unsigned char *p;
    if(!Nonce) return;
    p = Nonce;
    *(p+15)+=counter;
}


int rc6_demo()  //rc6_demo
{
    unsigned char *str = "rc6 encrypt test !";
    unsigned char *key = "1234567890ABCDEF";
    unsigned char *iv  = "1234567890ABCDEF";
    unsigned char arr_encrypt[64];
    unsigned char arr_decrypt[64];  //make sure have enough space
    struct rc6_ctx ctx = {
        .mode = RC6_MODE_CFB,
        .padding = RC6_PKCS7,
        .ctr_alg = default_ctr_alg
    };
    int ret;

    rc6_init(&ctx, key, 16, iv, NULL);
    rc6_encrypt(&ctx, str, strlen(str), arr_encrypt, &ret);
    rc6_decrypt(&ctx, arr_decrypt, &ret, arr_encrypt, ret);
    rc6_free(&ctx);

    printf("original: %s\n", str);
    printf("encrypt: %s\n", arr_encrypt);
    printf("decrypt: %s\n",arr_decrypt);

    return 0;
}